#!/usr/bin/env python3
"""
build_venv.py
在 Git 仓库根目录下创建 / 更新 .venv 虚拟环境，并确保 setuptools、build、pip 为最新版。
用法:
    python build_venv.py          # 若 .venv 已存在则复用，否则新建
    python build_venv.py --force  # 删除旧环境后重新创建
"""

from __future__ import annotations

import argparse
import asyncio
import subprocess
import sys
from pathlib import Path
from venv import EnvBuilder

REPO_ROOT = Path(__file__).resolve().parent
VENV_DIR = REPO_ROOT / ".venv"
PYPI_MIRROR = "https://mirrors.ustc.edu.cn/pypi/simple"  # 加速下载


class VenvBuilder(EnvBuilder):
    """扩展 EnvBuilder, 在环境创建后自动升级核心打包依赖。"""

    def __init__(self, *, upgrade_deps: bool = True, **kwargs) -> None:
        # Python 3.9+ 才支持 super() 无参写法
        super().__init__(**kwargs)
        self.upgrade_deps = upgrade_deps

    def post_setup(self, context) -> None:
        """环境中 setuptools、pip、build 强制升级到最新版。"""
        packages = ("pip", "setuptools", "build")
        python = str(Path(context.env_exe).resolve())
        cmd = [python, "-m", "pip", "install", "-i", PYPI_MIRROR, "-U", *packages]
        print(">>>", " ".join(cmd))
        subprocess.check_call(cmd)


##======================== 创建/更新 虚拟环境
async def create_or_update_venv(*, force: bool = False) -> None:
    if VENV_DIR.exists() and force:
        print(f"🗑 删除旧环境: {VENV_DIR}")
        import shutil

        shutil.rmtree(VENV_DIR)

    builder = VenvBuilder(
        system_site_packages=True,  # 允许使用系统包
        with_pip=True,  # 确保自带 pip
        upgrade_deps=True,  # 创建后升级依赖
    )
    print(f"🐍 创建 / 更新虚拟环境: {VENV_DIR}")
    builder.create(VENV_DIR)

    print("\n✅ 完成！激活环境使用: ")
    if sys.platform == "win32":
        print(f"    . {VENV_DIR}\\Scripts\\Activate.ps1")
    else:
        print(f"    source {VENV_DIR}/bin/activate")


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="为项目创建 / 更新 .venv 虚拟环境")
    parser.add_argument("-f", "--force", action="store_true", help="若环境已存在则先删除再重建")
    args = parser.parse_args(argv)

    try:
        asyncio.run(create_or_update_venv(force=args.force))
    except subprocess.CalledProcessError as e:
        print("❌ 安装依赖失败:", e, file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
